`timescale 1ns/1ps
module riscv_npu_top import riscv_pkg::*; (
  input  logic        clk,
  input  logic        rst_n,

  // JTAG (上电调试)
  input  logic        tck,
  input  logic        tms,
  input  logic        tdi,
  output logic        tdo,

  // APB Debug I/O (数据调试)
  input  logic [31:0] paddr,
  input  logic        psel,
  input  logic        penable,
  input  logic        pwrite,
  input  logic [31:0] pwdata,
  output logic [31:0] prdata,
  output logic        pready,
  output logic        pslverr,

  // AXI4-Lite Master 输出（连接到 SoC 总线）
  output logic [31:0] axi_awaddr,
  output logic        axi_awvalid,
  input  logic        axi_awready,
  output logic [31:0] axi_wdata,
  output logic [3:0]  axi_wstrb,
  output logic        axi_wvalid,
  input  logic        axi_wready,
  input  logic [1:0]  axi_bresp,
  input  logic        axi_bvalid,
  output logic        axi_bready,
  output logic [31:0] axi_araddr,
  output logic        axi_arvalid,
  input  logic        axi_arready,
  input  logic [31:0] axi_rdata,
  input  logic [1:0]  axi_rresp,
  input  logic        axi_rvalid,
  output logic        axi_rready
);

  // ==================================================================
  // 1) RISC-V Core (RV32IM)
  // ==================================================================
  logic        core_mem_req;
  logic        core_mem_we;
  logic [31:0] core_mem_addr;
  logic [31:0] core_mem_wdata;
  logic [31:0] core_mem_rdata;
  logic        core_mem_ack;

  rv32im_core u_core (
    .clk          (clk),
    .rst_n        (rst_n),
    .mem_req_o    (core_mem_req),
    .mem_we_o     (core_mem_we),
    .mem_addr_o   (core_mem_addr),
    .mem_wdata_o  (core_mem_wdata),
    .mem_rdata_i  (core_mem_rdata),
    .mem_ack_i    (core_mem_ack)
  );

  // ==================================================================
  // 2) Unified Memory (16 KB = 4K x 32) — 地址对齐修正
  // ==================================================================
  logic        mem_cs;
  logic        mem_we;
  logic [31:0] mem_addr;
  logic [31:0] mem_wdata;
  logic [31:0] mem_rdata;

  assign mem_cs    = (core_mem_addr[31:14] == 18'h0);        // 0x0000_0000 ~ 0x0000_3FFF
  assign mem_we    = core_mem_we & mem_cs;
  assign mem_addr  = {core_mem_addr[31:2], 2'b00};           // 强制 4-byte 对齐
  assign mem_wdata = core_mem_wdata;

  bram_4kx32 u_bram (
    .clk  (clk),
    .we   (mem_we),
    .addr (mem_addr[13:2]),    // 字地址 [11:0] → 4K
    .din  (mem_wdata),
    .dout (mem_rdata)
  );

  // ==================================================================
  // 3) NPU 加速器 (Dot-Product + ReLU) — AXI 直连
  // ==================================================================
  logic        npu_cs;
  logic        npu_we;
  logic [31:0] npu_addr;
  logic [31:0] npu_wdata;
  logic [31:0] npu_rdata;
  logic        npu_ack;

  assign npu_cs    = (core_mem_addr[31:28] == 4'h1);         // 0x1000_0000+
  assign npu_we    = core_mem_we & npu_cs;
  assign npu_addr  = core_mem_addr;
  assign npu_wdata = core_mem_wdata;

  npu_dotp_relu u_npu (
    .clk        (clk),
    .rst_n      (rst_n),
    .cs_i       (npu_cs),
    .we_i       (npu_we),
    .addr_i     (npu_addr[7:0]),
    .wdata_i    (npu_wdata),
    .rdata_o    (npu_rdata),
    .ack_o      (npu_ack),

    // AXI Master 接口（直通到顶层）
    .axi_awaddr (axi_awaddr),
    .axi_awvalid(axi_awvalid),
    .axi_awready(axi_awready),
    .axi_wdata  (axi_wdata),
    .axi_wstrb  (axi_wstrb),
    .axi_wvalid (axi_wvalid),
    .axi_wready (axi_wready),
    .axi_bresp  (axi_bresp),
    .axi_bvalid (axi_bvalid),
    .axi_bready (axi_bready),
    .axi_araddr (axi_araddr),
    .axi_arvalid(axi_arvalid),
    .axi_arready(axi_arready),
    .axi_rdata  (axi_rdata),
    .axi_rresp  (axi_rresp),
    .axi_rvalid (axi_rvalid),
    .axi_rready (axi_rready)
  );

  // ==================================================================
  // 4) 总线仲裁 & 响应
  // ==================================================================
  always_comb begin
    core_mem_rdata = 32'd0;
    core_mem_ack   = 1'b0;
    if (mem_cs) begin
      core_mem_rdata = mem_rdata;
      core_mem_ack   = 1'b1;
    end else if (npu_cs) begin
      core_mem_rdata = npu_rdata;
      core_mem_ack   = npu_ack;
    end
  end

  // ==================================================================
  // 5) JTAG TAP + 写寄存器支持
  // ==================================================================
  logic        jtag_ir_capture;
  logic [7:0]  jtag_ir;
  logic        jtag_dr_capture;
  logic [31:0] jtag_dr_in, jtag_dr_out;
  logic        jtag_update;

  jtag_tap u_jtag (
    .tck        (tck),
    .tms        (tms),
    .tdi        (tdi),
    .tdo        (tdo),
    .ir_capture (jtag_ir_capture),
    .ir_in      (jtag_ir),
    .dr_capture (jtag_dr_capture),
    .dr_in      (jtag_dr_in),
    .dr_out     (jtag_dr_out),
    .update_dr  (jtag_update)
  );

  // Debug CSR: PC + GPR[0..3]
  logic [31:0] dbg_csr [0:4];
  assign dbg_csr[0] = u_core.pc;
  assign dbg_csr[1] = u_core.gpr[0];
  assign dbg_csr[2] = u_core.gpr[1];
  assign dbg_csr[3] = u_core.gpr[2];
  assign dbg_csr[4] = u_core.gpr[3];

  // JTAG 读
  always_comb begin
    jtag_dr_in = 32'd0;
    if (jtag_ir >= 8'h01 && jtag_ir <= 8'h05)
      jtag_dr_in = dbg_csr[jtag_ir[2:0] - 1];
  end

  // JTAG 写 GPR（IR=0x10~0x13）
  always_ff @(posedge tck) begin
    if (jtag_update && jtag_ir >= 8'h10 && jtag_ir <= 8'h13) begin
      u_core.gpr[jtag_ir[1:0]] <= jtag_dr_out;
    end
  end

  // ==================================================================
  // 6) APB Debug Port — 完整桥接实现
  // ==================================================================
  logic        apb_mem_req;
  logic        apb_mem_we;
  logic [31:0] apb_mem_addr;
  logic [31:0] apb_mem_wdata;
  logic [31:0] apb_mem_rdata;
  logic        apb_mem_ack;

  apb_debug_port u_apb_dbg (
    .clk        (clk),
    .rst_n      (rst_n),
    .paddr_i    (paddr),
    .psel_i     (psel),
    .penable_i  (penable),
    .pwrite_i   (pwrite),
    .pwdata_i   (pwdata),
    .prdata_o   (prdata),
    .pready_o   (pready),
    .pslverr_o  (pslverr),

    // 桥接到统一地址空间
    .mem_req_o  (apb_mem_req),
    .mem_we_o   (apb_mem_we),
    .mem_addr_o (apb_mem_addr),
    .mem_wdata_o(apb_mem_wdata),
    .mem_rdata_i(apb_mem_rdata),
    .mem_ack_i  (apb_mem_ack)
  );

  // APB 请求仲裁（优先于 CPU）
  logic [31:0] final_mem_rdata;
  logic        final_mem_ack;

  always_comb begin
    final_mem_rdata = 32'd0;
    final_mem_ack   = 1'b0;

    if (apb_mem_req) begin
      final_mem_rdata = (apb_mem_addr[31:14] == 18'h0) ? mem_rdata :
                        (apb_mem_addr[31:28] == 4'h1)  ? npu_rdata : 32'd0;
      final_mem_ack   = 1'b1;
    end else if (core_mem_req) begin
      final_mem_rdata = core_mem_rdata;
      final_mem_ack   = core_mem_ack;
    end
  end

  assign apb_mem_rdata = final_mem_rdata;
  assign apb_mem_ack   = final_mem_ack;

  // 写路径：APB 和 CPU 互斥
  assign mem_we = apb_mem_req ? apb_mem_we :
                  core_mem_req ? core_mem_we : 1'b0;
  assign mem_addr  = apb_mem_req ? apb_mem_addr : core_mem_addr;
  assign mem_wdata = apb_mem_req ? apb_mem_wdata : core_mem_wdata;

  // NPU CSR 写（APB 直通）
  assign npu_we = apb_mem_req ? (apb_mem_we & (apb_mem_addr[31:28] == 4'h1)) :
                  core_mem_req ? npu_we : 1'b0;

endmodule